# @FileName  : main.py
# @Time      : 2025/2/6 7:18
# @Author    : LuZhaoHui
# @Software  : PyCharm
from file import *
from mainUI import Ui_MainWindow
from vars import *
from conf import *
from screen import *
from dt import *
from str import *
from api import *

import os, sys


class MainRun(QMainWindow, Ui_MainWindow):
    stdoutBak = ''
    stderrBak = ''
    # 路径
    dataPath = ''
    dataPathList = []
    curFileList = []
    curFile = ''
    resultPath = ''
    resultPathList = []
    # 模型
    process_reason_thread = None
    userInfo = {}
    isvLLM = True
    isOpenAI = False
    baseUrlRequest = ''
    baseUrlOpenAI = ''
    isApiKey = False
    apiKey = ''
    reqType = ''
    isStream = False
    reqStream = False
    curModel = ''
    modelList = []
    modelNameList = []
    curSystemPrompt = ''
    systemPromptList = []
    # 推理
    imagePage = 'A4'
    curDir = ''
    dirType = ''
    curPrompt = ''
    promptList = []
    # 数据
    isMysql = False
    mysqlHost = ''
    mysqlPort = 3306
    mysqlUser = ''
    mysqlPass = ''
    mysqlDatabase = ''
    msql = None
    #
    isReason = False
    start = False
    timer = None
    timer_count = 0

    def __init__(self, parent=None):
        super(MainRun, self).__init__(parent)
        save_log('start')
        self.setupUi(self)
        self.init_redirection()
        self.config = MainConf('main.ini')
        self.msql = None
        self.api = None
        self.init_var()
        self.init()
        self.init_ui()
        self.init_mysql()
        self.init_api()

    def init_redirection(self):
        self.stdoutBak = sys.stdout
        self.stderrBak = sys.stderr

    def reset_redirection(self):
        sys.stdout = self.stdoutBak
        sys.stderr = self.stderrBak

    def closeEvent(self, event, **kwargs):
        if not self.start and showBoxYes(self, "系统将退出，是否确认?  "):
            self.exit()
            event.accept()
        else:
            event.ignore()

    def exit(self):
        save_log('exit')
        if self.msql is not None:
            del self.msql
        if self.api is not None:
            del self.api
        if self.timer is not None:
            self.timer.stop()
            del self.timer
        self.reset_redirection()

    def init_var(self):
        # 路径
        self.dataPath = self.config.varConf['DataPath']
        self.dataPathList = []
        self.curFileList = []
        self.curFile = ''
        self.resultPath = self.config.varConf['ResultPath']
        self.resultPathList = []
        # 模型
        self.isvLLM = self.config.varModel['IsvLLM']
        self.isOpenAI = self.config.varModel['IsOpenAI']
        self.baseUrlRequest = self.config.varModel['BaseUrlRequest']
        self.baseUrlOpenAI = self.config.varModel['BaseUrlOpenAI']
        self.isApiKey = self.config.varModel['IsApiKey']
        self.apiKey = self.config.varModel['ApiKey']
        self.reqType = self.config.varModel['ReqType']
        self.isStream = self.config.varModel['IsStream']
        self.reqStream = self.config.varModel['ReqStream']
        self.curModel = self.config.varModel['CurModel']
        self.modelList = self.config.varModel['ModelList']
        self.modelNameList = self.config.varModel['ModelNameList']
        self.curSystemPrompt = self.config.varModel['CurSystemPrompt']
        self.systemPromptList = self.config.varModel['SystemPromptList']
        # 推理
        self.imagePage = self.config.varReason['ImagePage']
        self.curDir = self.config.varReason['CurDir']
        self.dirType = self.config.varReason['DirType']
        self.curPrompt = self.config.varReason['CurPrompt']
        self.promptList = self.config.varReason['PromptList']
        # 数据
        self.isMysql = self.config.varMysql['IsMysql']
        self.mysqlHost = self.config.varMysql['MysqlHost']
        self.mysqlPort = self.config.varMysql['MysqlPort']
        self.mysqlUser = self.config.varMysql['MysqlUser']
        self.mysqlPass = self.config.varMysql['MysqlPass']
        self.mysqlDatabase = self.config.varMysql['MysqlDatabase']
        self.msql = None
        #
        self.start = False

    def init(self):
        if self.config.varConf['WindowMax']:
            self.showMaximized()
        # self.setWindowFlags(Qt.WindowCloseButtonHint | Qt.WindowMinimizeButtonHint)
        self.setWindowIcon(QIcon(LOGO))
        self.setWindowTitle(' %s ver[%s]' % (self.config.varConf['Title'], VER))
        # 路径
        self.select_data_dir(self.dataPath)
        self.select_result_dir(self.resultPath)
        # 模型
        self.comboBoxCurModel.addItems(self.modelList)
        self.comboBoxCurModel.setCurrentIndex(self.curModel)
        self.comboBoxSystemPrompt.addItems(self.systemPromptList)
        self.comboBoxSystemPrompt.setCurrentIndex(self.curSystemPrompt)
        # 推理
        if self.isOpenAI:
            self.comboBoxReqType.addItem('openai')
        self.comboBoxReqType.setCurrentText(self.reqType)
        self.checkBoxReqStream.setEnabled(self.isStream)
        if self.isStream:
            self.checkBoxReqStream.setChecked(self.reqStream)
        self.comboBoxCurDir.setCurrentText(self.curDir)
        # URL,KEY
        self.lineEditApiKey.setText(self.apiKey)
        self.flush_url()
        # 数据
        self.checkBoxMysql.setChecked(self.isMysql)
        self.lineEditMysqlInfo.setText(
            '%s@%s:%d->%s' % (self.mysqlUser,
                              self.mysqlHost,
                              self.mysqlPort,
                              self.mysqlDatabase))
        # 定时器
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_time)
        self.timer.start(1000)
        # 提示
        self.flush_prompt_list()

    def init_ui(self):
        setComboBoxHeight(self.comboBoxCurModel, 18)
        setComboBoxHeight(self.comboBoxSystemPrompt, 18)
        setComboBoxHeight(self.comboBoxReqType, 18)
        setComboBoxHeight(self.comboBoxCurDir, 18)
        # 选择模型
        self.comboBoxCurModel.currentIndexChanged.connect(self.change_cur_model)
        # 选择提示
        self.comboBoxSystemPrompt.currentIndexChanged.connect(self.change_system_prompt)
        self.comboBoxSystemPrompt.lineEdit().editingFinished.connect(self.change_edit_system_prompt)
        # 请求类型
        self.comboBoxReqType.currentTextChanged.connect(self.change_req_type)
        # 当前目录
        self.comboBoxCurDir.currentTextChanged.connect(self.change_cur_dir)
        # 选择路径
        self.pushButtonSelectDataDir.clicked.connect(self.change_data_dir)
        self.pushButtonSelectResultDir.clicked.connect(self.change_result_dir)
        # 数据库，输出流
        self.checkBoxMysql.clicked.connect(self.change_mysql)
        if self.isStream:
            self.checkBoxReqStream.clicked.connect(self.change_req_stream)
        # 文件列表
        self.listWidgetCurFileList.itemClicked.connect(self.select_cur_file)
        self.listWidgetPrompt.itemClicked.connect(self.select_prompt)
        self.listWidgetPrompt.doubleClicked.connect(self.delete_prompt)
        # 推理
        self.pushButtonReason.clicked.connect(self.process_reason)
        self.pushButtonSave.clicked.connect(self.result_save)

    def init_mysql(self):
        if self.msql is not None:
            return

    def init_api(self):
        if self.api is not None:
            return
        self.api = modelApi(self.baseUrlRequest, self.baseUrlOpenAI, self.apiKey if self.isApiKey else None)
        self.api.setImagePage(self.imagePage)
        self.api.setIsvLLM(self.isvLLM)
        self.api.setReqType(self.reqType)
        self.api.setReqStream(self.reqStream)
        self.api.setSystemPrompt(self.systemPromptList[self.curSystemPrompt])
        self.api.setModel(self.modelNameList[self.curModel], not is_vl(self.modelList[self.curModel]))

    def change_mysql(self):
        # 是否数据库
        if self.isMysql != self.checkBoxMysql.isChecked():
            self.isMysql = self.checkBoxMysql.isChecked()
            self.config.write('MYSQL', 'IsMysql', self.isMysql)

    def change_req_stream(self):
        # 输出流
        if self.reqStream != self.checkBoxReqStream.isChecked():
            self.reqStream = self.checkBoxReqStream.isChecked()
            self.api.setReqStream(self.reqStream)
            self.config.write('MODEL', 'ReqStream', self.reqStream)

    def flush_url(self):
        if self.reqType == 'openai':
            self.lineEditBaseUrl.setText(self.baseUrlOpenAI)
        else:
            self.lineEditBaseUrl.setText(self.baseUrlRequest)

    def change_req_type(self, value):
        # 请求类型
        # print(value)
        if self.reqType != value:
            self.reqType = value
            self.api.setReqType(self.reqType)
            self.config.write('MODEL', 'ReqType', self.reqType)
            self.flush_url()

    def change_cur_model(self, value):
        # 当前模型
        if self.curModel != value:
            self.curModel = value
            self.api.setModel(self.modelNameList[self.curModel], not is_vl(self.modelList[self.curModel]))
            self.config.write('MODEL', 'CurModel', self.curModel)

    def change_system_prompt(self, value):
        # 当前提示
        if self.curSystemPrompt != value:
            self.curSystemPrompt = value
            self.api.setSystemPrompt(self.systemPromptList[self.curSystemPrompt])
            self.config.write('MODEL', 'CurSystemPrompt', self.curSystemPrompt)

    def change_edit_system_prompt(self):
        # 当前编辑提示
        text = self.comboBoxSystemPrompt.currentText()
        self.comboBoxSystemPrompt.setItemText(self.curSystemPrompt, text)
        self.systemPromptList.clear()
        s = ''
        for i in range(self.comboBoxSystemPrompt.count()):
            self.systemPromptList.append(self.comboBoxSystemPrompt.itemText(i))
            if i > 0:
                s += ','
            s += self.systemPromptList[i]
        self.config.write('MODEL', 'SystemPromptList', s)

    def change_cur_dir(self, value):
        # 当前目录
        self.curDir = value
        self.config.write(section='REASON', key='CurDir', value=value)
        path = path_join(self.dataPath, value)
        print(path)
        # 读取文件
        self.listWidgetCurFileList.clear()
        self.curFileList.clear()
        self.curFile = ''
        dirs = [path]
        while len(dirs):
            dir_path = dirs.pop()
            print('read dir [%s]' % dir_path)
            for d in os.listdir(dir_path):
                name = path_join(dir_path, d)
                if os.path.isdir(name):
                    dirs.append(name)
                    continue
                # 文件
                item_text = name
                self.listWidgetCurFileList.addItem(QListWidgetItem(item_text))
                self.curFileList.append(item_text)

    def select_cur_file(self):
        # 选择文件
        index = self.listWidgetCurFileList.selectedIndexes()[0].row()
        self.curFile = self.curFileList[index]
        showLabelImage(self.scrollAreaWidgetImage, self.labelImage, self.curFile, self=False)

    def select_data_dir(self, data_dir):
        # 数据目录
        self.dataPathList.clear()
        self.curFileList.clear()
        self.curFile = ''
        if os.path.exists(data_dir):
            self.dataPath = data_dir
            self.config.write('CONF', 'DataPath', self.dataPath)
            # 读取目录
            for d in os.listdir(self.dataPath):
                name = path_join(self.dataPath, d)
                if os.path.isdir(name):
                    self.dataPathList.append(d)
        else:
            self.dataPath = ''
        self.lineEditDataDir.setText(self.dataPath)
        self.comboBoxCurDir.blockSignals(True)
        self.comboBoxCurDir.clear()
        self.comboBoxCurDir.addItems(self.dataPathList)
        if len(self.dataPathList) > 0:
            if self.curDir not in self.dataPathList:
                self.curDir = self.dataPathList[0]
            self.change_cur_dir(self.curDir)
        self.comboBoxCurDir.blockSignals(False)

    def change_data_dir(self):
        # 选择数据目录
        path = QFileDialog.getExistingDirectory(self, " 选择图片主目录 ", self.dataPath)
        if os.path.exists(path):
            self.select_data_dir(path)

    def select_result_dir(self, result_dir):
        # 结果目录
        self.resultPathList.clear()
        if os.path.exists(result_dir):
            self.resultPath = result_dir
            self.config.write('CONF', 'ResultPath', self.resultPath)
            # 读取目录
            for d in os.listdir(self.resultPath):
                name = path_join(self.resultPath, d)
                if os.path.isdir(name):
                    self.resultPathList.append(d)
        else:
            self.resultPath = ''
        self.lineEditResultDir.setText(self.resultPath)

    def change_result_dir(self):
        # 选择结果目录
        path = QFileDialog.getExistingDirectory(self, " 选择结果目录 ", self.resultPath)
        if os.path.exists(path):
            self.selectResultDir(path)

    def flush_prompt_list(self):
        self.listWidgetPrompt.clear()
        self.listWidgetPrompt.addItems(self.promptList)

    def select_prompt(self):
        # 选择提示
        index = self.listWidgetPrompt.selectedIndexes()[0].row()
        self.curPrompt = self.listWidgetPrompt.item(index).text()
        self.textEditPrompt.setText(self.curPrompt)

    def delete_prompt(self):
        # 删除提示
        index = self.listWidgetPrompt.selectedIndexes()[0].row()
        text = self.listWidgetPrompt.item(index).text()
        print(text)
        if showBoxYes(self, '删除提示[%s]，是否确认?  ' % text):
            item = self.listWidgetPrompt.takeItem(index)
            if item:
                self.promptList.pop(index)
                self.save_prompt('')
                if self.curPrompt == text:
                    self.curPrompt = ''
                    self.textEditPrompt.clear()

    def save_prompt(self, text):
        # 保存提示
        if text != '':
            self.curPrompt = text
        if text == '' or self.curPrompt not in self.promptList:
            if text != '':
                self.promptList.insert(0, self.curPrompt)
                nums = self.config.varReason['PromptListNums']
                if nums and len(self.promptList) > nums:
                    self.promptList.pop()
            s = ''
            for i in range(len(self.promptList)):
                if i > 0:
                    s += ','
                s += self.promptList[i]
            self.config.write(section='REASON', key='PromptList', value=s)
            self.flush_prompt_list()

    def process_reason(self):
        # 推理
        text = self.textEditPrompt.toPlainText()
        if text == '':
            showBox(self, msg='请输入对话文本')
            return
        if is_vl(self.modelList[self.curModel]) and not os.path.isfile(self.curFile):
            showBox(self, msg='请选择图片')
            return
        if self.curPrompt != text:
            self.save_prompt(text)
        # if self.checkBoxRepeat:

        print('processReason')
        self.isReason = True
        self.timer_count = 0
        self.update_time()
        self.textEditResult.clear()
        self.pushButtonReason.setEnabled(False)
        self.process_reason_thread = ProcessReasonThread(self.api, self.curPrompt, self.curFile)
        self.process_reason_thread.update_item.connect(self.show_update_reason)
        self.process_reason_thread.start()

    def show_update_reason(self, list_msg):
        if len(list_msg):
            # 推理结果
            # result = {
            #     'id': '请求ID',
            #     'object': '返回结果类型',
            #     'created': 0,
            #     'model': '使用的推理模型',
            #     'content': '推理文本结果',
            #     'tokens': 0
            # }
            result = list_msg[0]
            self.textEditResult.setText(result)
            self.result_save()
            print(result)
            if self.config.varModel['IsGetUserInFo']:
                # 获取在线用户信息
                self.userInfo = self.api.get_user_info()
                if self.userInfo['status']:
                    self.labelUserInfo.setText(' 余额：%s  用户信息:%s %s' % (
                        self.userInfo['data']['balance'],
                        self.userInfo['data']['name'],
                        self.userInfo['data']['email']
                    ))
        # 结束
        self.isReason = False
        self.pushButtonReason.setEnabled(True)
        self.pushButtonReason.setText('推理')

    def result_save(self):
        result = self.textEditResult.toPlainText()
        # if result == '':
        #     return
        # 保存结果
        dir_name = path_join(self.resultPath, self.modelList[self.curModel])
        if not os.path.isdir(dir_name):
            os.mkdir(dir_name)

        msg = '%s 模型[%s] 类型[%s] 网址[%s] 对话[%s->%s] 结果如下：\n' % (
            now_datetime_str(),
            self.modelNameList[self.curModel],
            self.reqType,
            self.lineEditBaseUrl.text(),
            self.comboBoxSystemPrompt.currentText(),
            self.textEditPrompt.toPlainText())

        if not is_vl(self.modelList[self.curModel]):
            # 文本
            msg += result + '\n\n\n'
            name = path_join(self.resultPath, 'result.txt')
        else:
            msg += '图片[%s] \n %s \n\n\n' % (self.curFile, result)
            dir_name = path_join(dir_name, self.curDir)
            if not os.path.isdir(dir_name):
                os.mkdir(dir_name)
            file = os.path.basename(self.curFile).split('.')[0]
            # 图片
            name = path_join(dir_name, '%s.txt' % file)
        save_msg(msg, name)

    def update_time(self):
        if not self.isReason:
            return
        self.timer_count += 1
        self.pushButtonReason.setText('%d' % self.timer_count)


class ProcessReasonThread(QThread):
    # 处理推理线程
    update_item = pyqtSignal(list)

    def __init__(self, api, prompt, file, parent=None):
        super().__init__(parent)
        self.api = api
        self.prompt = prompt
        self.file = file

    def run(self):
        if self.api.isvLLM:
            print('%s\n %s %s' % (self.api.baseUrlRequest,
                                  self.api.curModel,
                                  self.prompt
                                  )
                  )
            result = self.api.reasonImageRequest_vLLM(self.file, self.prompt)
        else:
            print('%s\n %s\n %s\n %s %s %s\n %s %s' % (self.api.baseUrlRequest,
                                                       self.api.baseUrlOpenAI,
                                                       self.api.apiKey,
                                                       self.api.reqType, str(self.api.reqStream),
                                                       self.api.curModel,
                                                       self.api.systemPrompt,
                                                       self.prompt
                                                       )
                  )
            if self.api.reqType != 'openai':
                # request
                if not is_vl(self.api.curModel):
                    result = self.api.reasonTextRequest(self.prompt)
                else:
                    result = self.api.reasonImageRequest(self.file, self.prompt)
            else:
                # openai
                if not is_vl(self.api.curModel):
                    result = self.api.reasonTextOpenAI(self.prompt)
                else:
                    result = self.api.reasonImageOpenAI(self.file, self.prompt)
        if result is not None:
            self.update_item.emit([result])
        self.update_item.emit([])


if __name__ == '__main__':
    app = QApplication(sys.argv)

    main = MainRun()
    main.show()

    sys.exit(app.exec_())
